import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.OutputStream;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.util.Enumeration;
import java.util.zip.ZipEntry;
import java.util.zip.ZipException;
import java.util.zip.ZipFile;
import java.util.zip.ZipOutputStream;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;

/**
 * 通过密钥加密源目录文件到目标目录文件，命令行：securitkey.key publickey.cert bin bin-o [filter]
 * 
 * @author yiyongpeng
 * 
 */
public class EncryptClass {
	private static interface Filter {
		boolean filter(File src, File dist) throws IOException;
	}

	public static void main(String[] args) {
		String privateKeyFileName = args[0];// 私钥文件
		String publicKeyFileName = args[1];// 公钥文件，不存在自动生成
		File srcPath = new File(args[2]);// 源目录
		File distPath = new File(args[3]);// 目标目录
		final String filte = args.length > 4 ? args[4] : null;// 过滤表达式，不填表示加密所有
		if (filte != null) {
			System.err.println("filter: " + filte);
		}
		final Cipher cipher = cipher4encrypt(publicKeyFileName,
				privateKeyFileName);
		Filter filter = new Filter() {
			public boolean filter(File src, File dist) throws IOException {
				if (src.isDirectory()) {
					if (!dist.exists() && !dist.mkdir())
						throw new IOException("无法创建目录：" + dist);
					return true;
				} else if (src.isFile()) {
					if (filte != null && !src.getAbsolutePath().matches(filte))
						return false;
					if (!dist.exists() && !dist.createNewFile())
						throw new IOException("无法创建文件：" + dist);
					encrypt(src, dist, cipher);
					return false;
				} else {
					System.err.println("unkown file: " + src);
					return false;
				}
			}
		};
		try {
			// 目录或文件
			scanFile(srcPath, distPath, filter);
		} catch (IOException e1) {
			e1.printStackTrace();
		}
	}

	/**
	 * 扫描加密Jar包内部文件
	 * 
	 * @param srcJar
	 * @param distJar
	 * @throws IOException
	 * @throws IOException
	 */
	@SuppressWarnings("unchecked")
	private static void scanJar(ZipFile srcJar, File distJar, Cipher cipher) {
		ZipOutputStream out = null;
		Enumeration<ZipEntry> list = (Enumeration<ZipEntry>) srcJar.entries();
		try {
			if (!distJar.exists() && !distJar.createNewFile())
				throw new IOException("无法创建文件：" + distJar);
			out = new ZipOutputStream(new FileOutputStream(distJar));
			byte[] buf = new byte[1024];
			ByteArrayOutputStream bout = new ByteArrayOutputStream();
			for (; list.hasMoreElements();) {
				InputStream in = null;
				ZipEntry entry = list.nextElement();
				entry = new ZipEntry(entry.getName());
				out.putNextEntry(entry);
				if (!entry.isDirectory())
					try {
						in = srcJar.getInputStream(entry);
						int size;
						while ((size = in.read(buf)) > 0)
							bout.write(buf, 0, size);
						byte[] data = bout.toByteArray();
						bout.reset();
						// 只对class加密
						if (entry.getName().toLowerCase().endsWith(".class"))
							data = cipher.doFinal(data);
						out.write(data);
					} finally {
						if (in != null) {
							in.close();
							in = null;
						}
					}
				out.flush();
			}
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (IllegalBlockSizeException e) {
			e.printStackTrace();
		} catch (BadPaddingException e) {
			e.printStackTrace();
		} finally {
			if (out != null) {
				try {
					out.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}

	}

	/**
	 * 通过指定算法加密源文件到目标文件
	 * 
	 * @param src
	 * @param dist
	 * @param cipher
	 * @throws IOException
	 */
	protected static void encrypt(File src, File dist, Cipher cipher)
			throws IOException {
		// jar包内部文件加密
		if (src.getName().toLowerCase().endsWith(".jar"))
			try {
				ZipFile zip = new ZipFile(src);
				scanJar(zip, dist, cipher);
				return;
			} catch (ZipException e) {
				System.err.println("unjar:" + src);
			}
		// 普通文件加密整个文件
		InputStream in = null;
		OutputStream out = null;
		try {
			out = new BufferedOutputStream(new FileOutputStream(dist));
			in = new BufferedInputStream(new FileInputStream(src));
			ByteArrayOutputStream bout = new ByteArrayOutputStream();
			byte[] buf = new byte[1024];
			for (int size; (size = in.read(buf)) > 0;)
				bout.write(buf, 0, size);
			buf = bout.toByteArray();
			// 只对class加密
			if (src.getName().toLowerCase().endsWith(".class"))
				buf = cipher.doFinal(buf);
			out.write(buf);
			out.flush();
			System.out.println("encrypt:\n" + src + "\n" + dist);
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IllegalBlockSizeException e) {
			e.printStackTrace();
		} catch (BadPaddingException e) {
			e.printStackTrace();
		} finally {
			if (in != null) {
				in.close();
			}
			if (out != null) {
				out.close();
			}
		}
	}

	/**
	 * 扫描文件
	 * 
	 * @param src
	 * @param dist
	 * @param filte
	 * @param filter
	 * @throws IOException
	 */
	private static void scanFile(File src, File dist, Filter filter)
			throws IOException {
		if (!filter.filter(src, dist))
			return;
		if (src.isDirectory()) {
			File[] list = src.listFiles();
			if (list != null && list.length > 0)
				for (File file : list)
					scanFile(file, new File(dist, file.getName()), filter);
		}
	}

	/**
	 * 通过私钥生成公钥，并创建加密算法
	 * 
	 * @param publicKeyFileName
	 * @param privateKeyFileName
	 * @return
	 */
	public static Cipher cipher4encrypt(String publicKeyFileName,
			String privateKeyFileName) {
		ObjectInputStream priIn = null;
		OutputStream pubOut = null;
		try {
			// 反序列化私钥对象
			priIn = new ObjectInputStream(new BufferedInputStream(
					new FileInputStream(privateKeyFileName)));
			SecretKey priKey = (SecretKey) priIn.readObject();
			Key key = null;
			Cipher cipher = Cipher.getInstance(priKey.getAlgorithm());
			File file = new File(publicKeyFileName);
			if (!file.exists()) {
				// 生成公钥并用私钥包装输出到文件
				if (!file.createNewFile())
					throw new IOException("无法创建文件：" + file);
				KeyGenerator keyGen = KeyGenerator.getInstance(priKey
						.getAlgorithm());
				key = keyGen.generateKey();
				cipher.init(Cipher.WRAP_MODE, priKey);
				byte[] pubKey = cipher.wrap(key);
				pubOut = new BufferedOutputStream(new FileOutputStream(file));
				pubOut.write(pubKey);
				pubOut.flush();
				System.err.println("gen-pubkey: " + file + "\t"
						+ key.getAlgorithm());
			} else {
				// 读取并用私钥解包已存在的公钥文件
				InputStream pubIn = null;
				try {
					ByteArrayOutputStream out = new ByteArrayOutputStream();
					pubIn = new BufferedInputStream(new FileInputStream(file));
					byte[] buf = new byte[1024];
					for (int size; (size = pubIn.read(buf)) > 0;)
						out.write(buf, 0, size);
					buf = out.toByteArray();
					cipher.init(Cipher.UNWRAP_MODE, priKey);
					key = cipher.unwrap(buf, priKey.getAlgorithm(),
							Cipher.SECRET_KEY);
					System.err.println("load-pubkey: " + file + "\t"
							+ key.getAlgorithm());
				} finally {
					if (pubIn != null) {
						pubIn.close();
					}
				}
			}
			cipher.init(Cipher.ENCRYPT_MODE, key);
			return cipher;
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			e.printStackTrace();
		} catch (IllegalBlockSizeException e) {
			e.printStackTrace();
		} finally {
			if (priIn != null) {
				try {
					priIn.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
			if (pubOut != null) {
				try {
					pubOut.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}
		return null;
	}
}
